# -*- coding: utf-8 -*-

import tushare as ts
import datetime
import time
import pandas as pd

ts.set_token('1c1a98e5aff7015d61cf45b87830b6839da9b5313478c807e7c527f9')
pro = ts.pro_api('1c1a98e5aff7015d61cf45b87830b6839da9b5313478c807e7c527f9')
conStockData = pd.read_excel("2019年至今TOP5股票.xlsx",header = 1,index_col=0)
class Shuju:
    def __init__(self):
        pass

#对股票代码的格式进行补充
    def update_stockcode(stockcode):
        if stockcode[0] == '6':
            stockcode = stockcode + ".SH"
        else:
            stockcode = stockcode + ".SZ"
        return stockcode

#获得指定日期的收盘价
    def get_close_price(code, tradeday, conStockData):
        day_datetime = datetime.datetime.strptime(tradeday, '%Y%m%d')
        tradeday = day_datetime.strftime("%Y-%m-%d")
        code = code[:-3]
        return conStockData.loc[tradeday, code]

#获得下一个交易日的日期
    def get_next_trading(lastday):
        lastday = datetime.datetime.strptime(lastday, '%Y-%m-%d')
        lastday = lastday + datetime.timedelta(days=1)
        lastday_str = lastday.strftime("%Y%m%d")
        while not pro.query('trade_cal', start_date=lastday_str, end_date=lastday_str, fields="is_open").iloc[0, 0]:
            lastday = lastday + datetime.timedelta(days=1)
            lastday_str = lastday.strftime("%Y%m%d")
            time.sleep(0.12)
        lastday_str = lastday.strftime("%Y-%m-%d")
        return lastday_str

